package com.scala.learn.sparkUDF

import org.apache.orc.impl.TreeReaderFactory.LongTreeReader
import org.apache.spark.sql.Row
import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types._

/**
  * @Copyright: Shanghai Definesys Company.All rights reserved.
  * @Description:
  * @author: chuhaitao
  * @since: 2019/3/16 22:48
  * @history:
  *          1.2019/3/16 created by chuhaitao
  */
object Geo {

}

/**
  * 定义计算多个数的集合平均数
  * 例如：（1*2*3） 的1/3次方
  **/
class GeometricMean extends UserDefinedAggregateFunction {

  //输入的类型
  override def inputSchema: StructType = StructType(List(StructField("value", DoubleType)))

  //中间结果的Schema
  override def bufferSchema: StructType = StructType(
    List(
      StructField("count", LongType),
      StructField("product", DoubleType)
    ))

  //返回值类型
  override def dataType: DataType = DoubleType

  //用来标记正对给定的一组数组，函数是总生成相同的结果
  override def deterministic: Boolean = true

  //对聚合运算中间结果的初始化
  override def initialize(buffer: MutableAggregationBuffer): Unit = {

    buffer(0) = 0L
    buffer(1) = 1.0
  }

  //每一条数据都执行的操作
  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    buffer(0) = buffer.getAs[Long](0) + 1
    buffer(1) = buffer.getAs[Double](1) + input.getAs[Double](0)

  }

  //负责合并两个聚合运算的buffer，在存储到MutableAggregationBuffer中
  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {

    buffer1(0) = buffer1.getAs[Long](0) + buffer2.getAs[Long](0)
    buffer1(1) = buffer1.getAs[Double](1) * buffer2.getAs[Double](1)
  }

  //完成对聚合buffer的值的运算
  override def evaluate(buffer: Row): Any = {
    math.pow(buffer.getAs[Double](1), 1.toDouble / buffer.getLong(0))

  }
}
